import os

import numpy as np
import faiss
import base_util
import embeding_model


class FaissVectorDatabaseManager:
    def __init__(self,list_no=8, nlist=168):
        """
         list_no 一般是cpu或Gpu的整数倍 方便并行计算
        :param list_no:
        :param nlist:
        """
        self.indices = {}
        self.list_no=list_no
        self.nlist=nlist

        self.indexs_path=f"{base_util.current_project_dir}/indexs"
        filenames = os.listdir(self.indexs_path)
        for  f in filenames:
            if f.endswith(".index"):
                index_path=f"{self.indexs_path}/{f}"
                index_name=f.replace(".index","")
                print("加载index",index_name)
                self.indices[index_name] = faiss.read_index(index_path)


    def create_index(self, index_name, dimension):
        """
        创建一个新的索引，用于存储特定类别的向量。
        :param index_name: 索引的名称，类似于表名。
        :param vectors: 一个形状为 (num_vectors, dimension) 的 NumPy 数组。
        """
        if index_name in self.indices:
            raise ValueError(f"Index '{index_name}' already exists.")

        # 创建IndexIVFFlat索引
        base_index = faiss.IndexFlatL2(dimension)  # 基本索引
        index_id_map = faiss.IndexIDMap(base_index)
        self.indices[index_name] = index_id_map
        # self.indices[index_name].add(vectors)
        print(f"Index '{index_name}' created successfully.")

    def add_vectors_to_index(self, index_name, vectors,ids):
        """
        向指定索引添加向量。
        :param index_name: 索引的名称。
        :param vectors: 一个形状为 (num_vectors, dimension) 的 NumPy 数组。
        """
        if index_name not in self.indices:
            raise ValueError(f"Index '{index_name}' does not exist.")
        # self.indices[index_name].add(vectors)
        self.indices[index_name].add_with_ids(vectors,ids)
        print(f"Vectors added to index '{index_name}' successfully.")

    def search(self, index_name, query_vector, k=10):
        """
        在指定索引中执行相似度检索，并规范化返回的距离值。
        :param index_name: 索引的名称。
        :param query_vector: 一个形状为 (1, dimension) 的 NumPy 数组，表示查询向量。
        :param k: 要返回的最近邻的数量。
        :return: 返回最近邻的索引和规范化后的距离。
        """
        if index_name not in self.indices:
            raise ValueError(f"Index '{index_name}' does not exist.")

        # D, I = index.search(xq, k)

        distances, indices = self.indices[index_name].search(query_vector.reshape(1, -1), k)
        # # 规范化距离值
        normalized_distances = self.normalize_distances(distances)

        return indices[0], normalized_distances


    def normalize_distances(self, distances):
        """
        将距离值规范化到 0 到 1 之间。
        :param distances: 原始距离值列表。
        :return: 规范化后的距离值列表。
        """
        # 处理无穷大值
        finite_distances = np.where(np.isfinite(distances), distances, np.finfo(np.float32).max)
        # 规范化到 [0, 1]
        normalized = (finite_distances - np.min(finite_distances)) / (
                    np.max(finite_distances) - np.min(finite_distances))
        return normalized.tolist()


    def save(self,index_name):

        if index_name not in self.indices:
            raise ValueError(f"Index '{index_name}' does not exist.")

        base_index = self.indices[index_name].index
        # "./indexs/demo.index"
        index_save_path = f"{self.indexs_path}/{index_name}.index"
        faiss.write_index(base_index, index_save_path)




# 使用示例
if __name__ == "__main__":
    # 向量库的用法
    db_manager = FaissVectorDatabaseManager()
    # 首先创建一个BCEmbedding模型 (目前网易的模型  号称最好的embedding模型)
    emb_model = embeding_model.BCEmbedding()
    #  将 测试 字符串  转换成embeding
    e=emb_model.embed("测试")
    dim=e.shape[1]
    # 如果用相同的embeding模型的话 dim 可以写死
    print(dim)
    db_manager.create_index("dataset1", dim)

    # 生成10个字符串的向量 并存到数据库里
    # 这里的id 我取的uuid  后期将替换成mongodb的 _id
    #真实场景下就是 文本分割以后 将每个短文本存到mongodb中。  然后每个短文本转换成向量，  接着将向量和mongodb的_id 一起存到向量数据库中。
    for i in range(0,10):
        # np.random.rand(100, 128)
        n=emb_model.embed(str(i))
        id=base_util.get_uuid()
        print(str(i),id)
        db_manager.add_vectors_to_index("dataset1", n,id)

    c=emb_model.embed("测")
    db_manager.add_vectors_to_index("dataset1", c, 1233)
    db_manager.add_vectors_to_index("dataset1", e, 123)

    query_vector =emb_model.embed("测试")

    # 进行向量查询检索
    k = 5
    top_k_ids, distances = db_manager.search("dataset1", query_vector, k)
    print(f"Top-k ids: {top_k_ids}, Distances: {distances}")

    # db_manager.save('dataset1' )


    # db_manager = FaissVectorDatabaseManager()
    #
    # from embeding_model import  BgeEmbedding
    #
    # embedding_model = BgeEmbedding()
    #
    #
    # emb=embedding_model.embed("测试数据xxxx")
    #
    # print(emb)
    # # 创建两个索引，分别用于两个不同的数据集
    # db_manager.create_index("dataset1", emb)
    # # db_manager.create_index("dataset2", np.random.rand(50, 128))
    # #
    # db_manager.add_vectors_to_index('dataset1',emb)
    #
    # emb2 = embedding_model.embed("测试数据xxxx2")
    # db_manager.add_vectors_to_index('dataset1', emb2)
    # # # 在 dataset1 索引中搜索
    # query_vector = embedding_model.embed("测试数据xxxx")
    # k = 5
    # top_k_indices, distances = db_manager.search("dataset1", query_vector, k)
    # print(f"Top-k indices for dataset1: {top_k_indices}, 距离: {distances}")
    # #
    # # # 在 dataset2 索引中搜索
    # # top_k_indices, distances = db_manager.search("dataset2", query_vector, k)
    # # print(f"Top-k indices for dataset2: {top_k_indices}, Distances: {distances}")